
import os
import glob
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import classification_report

# —— 用户可选模型 ——
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
from catboost import CatBoostClassifier

available_models = {
    'RandomForest': RandomForestClassifier(n_estimators=100, random_state=42),
    'LogisticRegression': LogisticRegression(max_iter=1000),
    'XGBoost': XGBClassifier(use_label_encoder=False, eval_metric='logloss'),
    'LightGBM': LGBMClassifier(),
    'CatBoost': CatBoostClassifier(verbose=0)
}
# 选择要运行的模型
models_to_run = ['RandomForest', 'LogisticRegression', 'CatBoost']

# —— 用户可选特征组合 ——
feature_sets = {
    'core5':   ['sdnn','rmssd','hf','lf','sampen'],
    'time2':   ['sdnn','rmssd'],
    'freq2':   ['hf','lf'],
    'mixed3':  ['rmssd','hf','sampen']
}

# —— 年龄分组函数 ——
def assign_age_group_scheme1(age):
    if age < 30:      return 'G1(<30)'
    if age < 40:      return 'G2(30-39)'
    if age < 50:      return 'G3(40-49)'
    return 'G4(>=50)'

def assign_age_group_scheme2(age):
    if age < 25:      return 'G1(<25)'
    if age < 50:      return 'G2(25-49)'
    if age < 75:      return 'G3(50-74)'
    return 'G4(>=75)'

# —— 构建第3-8 window 特征平均 ——
def build_features(folder, min_w=10, s=2, e=8):
    recs=[]
    for fp in glob.glob(os.path.join(folder,'*_rrfeature.csv')):
        pid=os.path.basename(fp).replace('_rrfeature.csv','')
        df=pd.read_csv(fp)
        w=df[df['label']=='window']
        if len(w)<min_w: continue
        sub=w.iloc[s:e]
        rec={'patient_id':pid}
        for c in ['sdnn','rmssd','hf','lf','sampen']:
            vals=pd.to_numeric(sub[c],errors='coerce')
            rec[c]=vals.mean() if not vals.dropna().empty else np.nan
        recs.append(rec)
    return pd.DataFrame(recs)

# —— 读取并打标签 ——
def load_and_label(roster_csv, scheme):
    df=pd.read_csv(roster_csv, dtype=str)
    df['patient_id']=df['userid'].str.strip()+'_'+df['recordid'].str.strip()
    df['age']=pd.to_numeric(df['age'],errors='coerce')
    if scheme=='scheme1':
        df['age_group']=df['age'].apply(assign_age_group_scheme1)
    else:
        df['age_group']=df['age'].apply(assign_age_group_scheme2)
    return df[['patient_id','age_group']].dropna()

# —— 主流程：同时跑两种分组方案 ——
def main(feature_folder, roster_csv):
    feats_df = build_features(feature_folder)
    for scheme in ['scheme1', 'scheme2']:
        label_df = load_and_label(roster_csv, scheme)
        data = feats_df.merge(label_df, on='patient_id', how='inner').dropna()
        print(f"\n共 {len(data)} 个样本，使用分组方案: {scheme}\n")
        le = LabelEncoder()
        data['y'] = le.fit_transform(data['age_group'])
        for fs_name, feat_cols in feature_sets.items():
            X = data[feat_cols].replace([np.inf,-np.inf],np.nan).dropna()
            y = data.loc[X.index, 'y']
            if X.empty:
                print(f"⚠️ 特征组 {fs_name} 无有效样本，跳过")
                continue
            Xtr, Xte, ytr, yte = train_test_split(
                X.astype(float), y, test_size=0.3,
                random_state=42, stratify=y
            )
            print(f"--- 分组方案 {scheme}，特征组: {fs_name} ---")
            for name in models_to_run:
                clf = available_models[name]
                clf.fit(Xtr, ytr)
                yp = clf.predict(Xte)
                print(f"\n模型: {name}")
                print(classification_report(yte, yp, target_names=le.classes_, zero_division=0))

if __name__=='__main__':
    feature_folder = '/database/home/duansizhang/hrv_predict/data/rr_featrue/'
    roster_csv     = '/database/private/mgcdb/info_xml_raw_match_1.csv'
    main(feature_folder, roster_csv)

